"""
均线上穿下穿算法验证工具 (简化版)
功能：验证趋势追踪止损中使用的均线上穿下穿算法
输出：生成一个列表，标记每个K线周期的上穿、下穿或无信号状态
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqsdk import TqApi, TqAuth, TqKq
import time

# MyTT模块中的函数已在下方直接定义，无需额外导入MyTT库
# 这些是从MyTT库中提取的常用技术指标函数
def REF(Series, N=1):
    """
    向前引用N周期的数据
    """
    return pd.Series(Series).shift(N)

def EMA(Series, N=10):
    """
    指数移动平均线
    """
    return pd.Series(Series).ewm(span=N, adjust=False).mean()

def SUM(Series, N=20):
    """
    求和
    """
    return pd.Series(Series).rolling(N).sum()

def calculate_trading_signals(klines, data_length=500):
    """
    计算交易信号
    :param klines: K线数据
    :param data_length: 数据长度
    :return: 包含信号的DataFrame
    """
    # 确保数据长度足够
    if len(klines) < data_length:
        print(f"警告: K线数据长度不足 {data_length}，实际长度为 {len(klines)}")
    
    # 转换为numpy数组以提高计算效率
    close = np.array(klines.close)
    low = np.array(klines.low)
    open_price = np.array(klines.open)
    high = np.array(klines.high)
    
    # 计算Q值
    Q = (3 * close + low + open_price + high) / 6
    
    # 计算trading_line
    terms = [
        26 * Q,
        25 * REF(Q, 1),
        24 * REF(Q, 2),
        23 * REF(Q, 3),
        22 * REF(Q, 4),
        21 * REF(Q, 5),
        20 * REF(Q, 6),
        19 * REF(Q, 7),
        18 * REF(Q, 8),
        17 * REF(Q, 9),
        16 * REF(Q, 10),
        15 * REF(Q, 11),
        14 * REF(Q, 12),
        13 * REF(Q, 13),
        12 * REF(Q, 14),
        11 * REF(Q, 15),
        10 * REF(Q, 16),
        9 * REF(Q, 17),
        8 * REF(Q, 18),
        7 * REF(Q, 19),
        6 * REF(Q, 20),
        5 * REF(Q, 21),
        4 * REF(Q, 22),
        3 * REF(Q, 23),
        2 * REF(Q, 24),
        REF(Q, 25)
    ]
    
    trading_line = sum(terms) / 351
    short_line = EMA(trading_line, 7)
    
    # 创建结果DataFrame
    result = pd.DataFrame({
        'datetime': klines.datetime,
        'close': close,
        'trading_line': trading_line,
        'short_line': short_line
    })
    
    # 计算信号
    result['signal'] = np.nan  # 默认为NaN
    
    # 计算上穿下穿信号
    for i in range(2, len(result)):
        # 上穿信号: 前一周期trading_line < short_line，当前周期trading_line > short_line
        if (result.trading_line.iloc[i-1] < result.short_line.iloc[i-1] and 
            result.trading_line.iloc[i] > result.short_line.iloc[i]):
            result.loc[result.index[i], 'signal'] = '上穿'
        
        # 下穿信号: 前一周期trading_line > short_line，当前周期trading_line < short_line
        elif (result.trading_line.iloc[i-1] > result.short_line.iloc[i-1] and 
              result.trading_line.iloc[i] < result.short_line.iloc[i]):
            result.loc[result.index[i], 'signal'] = '下穿'
    
    return result

# 主函数
print("开始验证均线上穿下穿算法...")

try:
    # 创建API连接
    api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"))
    
    # 设置默认合约和周期
    symbol = "SHFE.au2412"  # 上期所黄金
    period = 86400  # 日线
    
    print(f"正在获取 {symbol} 的日线数据...")
    klines = api.get_kline_serial(symbol, period, data_length=500)
    
    # 计算信号
    result = calculate_trading_signals(klines)
    
    # 输出结果
    print("\n计算完成! 结果摘要:")
    print(f"总K线数量: {len(result)}")
    print(f"上穿信号数量: {result['signal'].value_counts().get('上穿', 0)}")
    print(f"下穿信号数量: {result['signal'].value_counts().get('下穿', 0)}")
    
    # 显示最近的10个信号
    recent_signals = result[result['signal'].notna()].tail(10)
    if not recent_signals.empty:
        print("\n最近的10个信号:")
        for idx, row in recent_signals.iterrows():
            date_str = pd.to_datetime(row['datetime'], unit='s').strftime('%Y-%m-%d %H:%M:%S')
            print(f"{date_str}: {row['signal']} (trading_line: {row['trading_line']:.2f}, short_line: {row['short_line']:.2f})")
    
    # 保存结果到CSV
    csv_filename = f"{symbol.replace('.', '_')}_日线_信号.csv"
    result.to_csv(csv_filename, index=False)
    print(f"\n结果已保存到文件: {csv_filename}")
    
    # 绘制图表
    plt.figure(figsize=(15, 10))
    
    # 绘制收盘价
    plt.subplot(2, 1, 1)
    plt.plot(result.index, result['close'], label='收盘价', color='black', alpha=0.6)
    plt.title(f"{symbol} 日线收盘价")
    plt.legend()
    plt.grid(True)
    
    # 绘制均线和信号
    plt.subplot(2, 1, 2)
    plt.plot(result.index, result['trading_line'], label='Trading Line', color='blue')
    plt.plot(result.index, result['short_line'], label='Short Line', color='red')
    
    # 标记上穿和下穿点
    up_cross = result[result['signal'] == '上穿']
    down_cross = result[result['signal'] == '下穿']
    
    plt.scatter(up_cross.index, up_cross['trading_line'], color='green', marker='^', s=100, label='上穿')
    plt.scatter(down_cross.index, down_cross['trading_line'], color='red', marker='v', s=100, label='下穿')
    
    plt.title(f"{symbol} 日线均线和信号")
    plt.legend()
    plt.grid(True)
    
    # 保存图表
    plt_filename = f"{symbol.replace('.', '_')}_日线_图表.png"
    plt.savefig(plt_filename)
    print(f"图表已保存到文件: {plt_filename}")
    
    # 显示图表
    plt.tight_layout()
    plt.show()
    
except Exception as e:
    print(f"发生错误: {e}")
finally:
    if 'api' in locals():
        api.close()
        print("API连接已关闭")